from datetime import datetime
import commands
import psycopg2
import hashlib
import locale
import gzip
import time
import sys
import re
import os

import utils


def downloadURL():
	'''
	cp data from urlcrawl machine and unzip those files.
	'''
	cmdrsync = 'rsync -avz urlcrawl:/space/malicious-urls/ /home/intern/chromium/webcrawl-malware/src/data/MALURLs/'
	# cmdcp = 'cp /space/malicious-urls/ ../../data/TempData/MALURLs/'
	cmdunzip = 'gunzip -k -f --suffix .gzip ../../data/TempData/MALURLs/*.gzip'
	# print cmdrsync + '\n' + cmdunzip
	print cmdrsync + '\n' + cmdunzip
	print commands.getoutput(cmdrsync + '; ' + cmdunzip)


def containsAnyStr(tosearch, strlist):
	for s in strlist:
		if s in tosearch:
			return True
	return False


def formatToURL(infile, outfile, thedate, allhash, alltype):
	print infile
	print outfile

	# set input and output
	inf = gzip.open(infile, 'rb')
	inf = inf.read().split('\n')
	outf = open(outfile, 'w')
	first = True

	# the malicious type of urls focus
	type_to_focus = ['Blacole', 'Exploit', 'Obfuscator', 'JS', 'Malicious']

	for line in inf:
		line = line.split(',')
		if len(line) != 6:
			continue
		link = line[0]
		domain = line[2]
		real_type = line[4]
		real_type = real_type.replace('\'', '')
		
		'''
		If the real_type is already mapped, assign the label directly
		Else if the real_type is what we want, map it and assign label
		'''
		if real_type in alltype:
			None
		elif containsAnyStr(real_type, type_to_focus):
			# real_type, M1
			alltype[real_type] = 'M' + str(len(alltype) + 1)
		else:
			continue

		label = alltype[real_type]
		hexhash = hashlib.sha256(link).hexdigest()
		# Avoid duplicate keys
		if hexhash in allhash:
			continue
		else:
			allhash[hexhash] = 1
		# Format for psql input
		if first:
			first = False
			outf.write('(\'' + hexhash + '\',\'' + link + '\',\'' + domain + '\',\'' + label + '\',\'' + real_type + '\',\'' + thedate + '\')')
		else:
			outf.write(',\n(\'' + hexhash + '\',\'' + link + '\',\'' + domain + '\',\'' + label + '\',\'' + real_type + '\',\'' + thedate + '\')')

		# outf.write(hexhash + ',' + link + ',' + domain + ',' + label + ',' + real_type + ',' + thedate + '\n')



def malurl(allhash, alltype, basedir):
	locale.setlocale( locale.LC_ALL, 'en_US.UTF-8' )

	# downloadURL()
	malurldir = "/space/malicious-urls/"
	cmdlsinput = 'ls ' + malurldir + 'www.malwareurl.com-*'
	inputfiles = commands.getoutput(cmdlsinput).split('\n')

	for inputfile in inputfiles:
		try:
			date = inputfile.split('/')[-1].split('-')[1:]
			date[2] = date[2].split('.')[0]		# get rid of the .gzip suffix
			date = datetime(int(date[0]), int(date[1]), int(date[2]))
		except:
			print 'Error: bad filename', inputfile
			continue
		thedate = date.strftime('%Y-%m-%d')
		outputfile = basedir + 'URLs/malurl-' + thedate
		if not os.path.exists(outputfile):
			print thedate
			formatToURL(inputfile, outputfile, thedate, allhash, alltype)

		# After each file is processed, save current state
		savestate(allhash, alltype)


def isHTML(filename):
	cmd = "file {0} | grep 'HTML'".format(filename)
	if commands.getoutput(cmd):
		return True
	return False


def formatFile(urlfile, outfile, indir, mapfile='../../data/JSTerry/mapfile'):
	# Used to deal with Terry's examples
	# get the filenames and format to formatToURL input
	urlf = open(urlfile, 'r')
	

	outf = open(outfile, 'w')
	mapf = open(mapfile, 'w')
	first = True
	for url in urlf:
		url = url[:-1]
		# return whether the file is html or not
		filename = url.split('/')[-1]
		fileToJudge = indir + url.split('/')[-1]
		if not isHTML(fileToJudge):
			url = 'ruian.gtisc.gatech.edu/Terry/showjs.php?jsname={0}'.format(filename)

		domain = 'ruian.gtisc.gatech.edu'
		label = 'MTerry'
		real_type = 'javascript exploit'
		hex_hash = hashlib.sha256(url).hexdigest()
		if first:
			first = False
			outf.write('(\'' + hex_hash + '\',\'' + url + '\',\'' + domain + '\',\'' + label + '\',\'' + real_type + '\')')
		else:
			outf.write(',\n(\'' + hex_hash + '\',\'' + url + '\',\'' + domain + '\',\'' + label + '\',\'' + real_type + '\')')

		mapf.write(hex_hash + ',' + filename + '\n')
		
	outf.close()
	mapf.close()



def loadstate():
	basedir = '../../data/TempData/LOGs/'
	# load the available hash from file
	infile = basedir + 'allhash'
	allhash = utils.DictFile().file2dict(infile, 'None')
	# load the available types from file
	infile = basedir + 'alltype'
	alltype = utils.DictFile().file2dict(infile, 'str')
	return allhash, alltype


def savestate(allhash, alltype):
	basedir = '../../data/TempData/LOGs/'
	# save the available hash to file
	outfile = basedir + 'allhash'
	utils.DictFile().dict2file(allhash, outfile)
	# save the available types to file
	outfile = basedir + 'alltype'
	utils.DictFile().dict2file(alltype, outfile)


def main():
	basedir = '../../data/TempData/'
	allhash, alltype = loadstate()
	while True:
		malurl(allhash, alltype, basedir)
		savestate(allhash, alltype)
		print 'Sleeping'
		time.sleep(60 * 60 * 24)


if __name__=="__main__":
	main()
